import pymysql
from pymysql.cursors import DictCursor
import os
from datetime import datetime
import sys
from pathlib import Path

# === 配置项目路径 ===
project_root = str(Path(__file__).resolve().parent.parent)
sys.path.append(project_root)

from src.datasets.veri776_dataset import Veri776Dataset


class DatabaseManager:
    def __init__(self, host='localhost', user='root', password='root', db='vehicle_reid'):
        self.connection = pymysql.connect(
            host=host,
            user=user,
            password=password,
            db=db,
            charset='utf8',
            cursorclass=DictCursor,
            autocommit=False
        )
        self.create_tables()

    def create_tables(self):
        """创建车辆表和三个模式的图片表"""
        with self.connection.cursor() as cursor:
            # 车辆表
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS vehicles (
                    vehicle_id INT PRIMARY KEY COMMENT '车辆唯一ID',
                    vehicle_name VARCHAR(100) COMMENT '车辆名称',
                    color VARCHAR(50) COMMENT '车辆颜色',
                    type VARCHAR(50) COMMENT '车辆类型'
                ) DEFAULT CHARSET=utf8;
            """)
            
            # 训练集图片表
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS images_train (
                    id INT AUTO_INCREMENT PRIMARY KEY COMMENT '自增ID',
                    vehicle_id INT COMMENT '关联的车辆ID',
                    image_path VARCHAR(255) NOT NULL COMMENT '图片路径（相对于datas/VeRi-776/）',
                    camera_id INT NOT NULL COMMENT '摄像头ID',
                    capture_time DATETIME COMMENT '拍摄时间',
                    FOREIGN KEY (vehicle_id) REFERENCES vehicles(vehicle_id)
                ) DEFAULT CHARSET=utf8;
            """)
            
            # 测试集图片表
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS images_test (
                    id INT AUTO_INCREMENT PRIMARY KEY COMMENT '自增ID',
                    vehicle_id INT COMMENT '关联的车辆ID',
                    image_path VARCHAR(255) NOT NULL COMMENT '图片路径（相对于datas/VeRi-776/）',
                    camera_id INT NOT NULL COMMENT '摄像头ID',
                    capture_time DATETIME COMMENT '拍摄时间',
                    FOREIGN KEY (vehicle_id) REFERENCES vehicles(vehicle_id)
                ) DEFAULT CHARSET=utf8;
            """)
            
            # 查询集图片表
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS images_query (
                    id INT AUTO_INCREMENT PRIMARY KEY COMMENT '自增ID',
                    vehicle_id INT COMMENT '关联的车辆ID',
                    image_path VARCHAR(255) NOT NULL COMMENT '图片路径（相对于datas/VeRi-776/）',
                    camera_id INT NOT NULL COMMENT '摄像头ID',
                    capture_time DATETIME COMMENT '拍摄时间',
                    FOREIGN KEY (vehicle_id) REFERENCES vehicles(vehicle_id)
                ) DEFAULT CHARSET=utf8;
            """)
        self.connection.commit()

    def import_veri776_dataset(self, dataset_path):
        """按模式分表导入数据集（修复外键约束问题）"""
        mode_tables = {
            "train": "images_train",
            "test": "images_test",
            "query": "images_query"
        }
        
        # 第一步：收集所有车辆信息（去重）
        vehicles = {}
        for mode in ["train", "test", "query"]:
            dataset = Veri776Dataset(mode=mode, root_dir=dataset_path)
            # 转换为相对路径并收集数据
            for item in dataset.data:
                img_path_abs = item[0]
                img_path_rel = os.path.relpath(img_path_abs, start=dataset_path)
                vehicle_id = item[1]
                # 收集车辆信息（仅首次出现时）
                if vehicle_id not in vehicles:
                    vehicles[vehicle_id] = {
                        "vehicle_id": vehicle_id,
                        "vehicle_name": f"vehicle_{vehicle_id}",
                        "color": "unknown",
                        "type": "unknown"
                    }
        
        # 第二步：先插入车辆表（关键修复点）
        self._insert_vehicles(vehicles)
        print(f"成功插入 {len(vehicles)} 条车辆记录")
        
        # 第三步：插入各模式的图片数据
        for mode in ["train", "test", "query"]:
            dataset = Veri776Dataset(mode=mode, root_dir=dataset_path)
            table_name = mode_tables[mode]
            
            # 转换为相对路径
            relative_data = []
            for item in dataset.data:
                img_path_abs = item[0]
                img_path_rel = os.path.relpath(img_path_abs, start=dataset_path)
                relative_data.append((img_path_rel, item[1], item[2]))  # (相对路径, 车辆ID, 摄像头ID)
            
            # 插入当前模式的图片数据
            self._insert_images(table_name, relative_data)

    def _insert_vehicles(self, vehicles):
        """插入车辆表"""
        with self.connection.cursor() as cursor:
            vehicle_sql = """
                INSERT IGNORE INTO vehicles (vehicle_id, vehicle_name, color, type) 
                VALUES (%s, %s, %s, %s)
            """
            cursor.executemany(vehicle_sql, [
                (v["vehicle_id"], v["vehicle_name"], v["color"], v["type"])
                for v in vehicles.values()
            ])
        self.connection.commit()

    def _insert_images(self, table_name, image_data):
        """插入指定表的图片数据"""
        with self.connection.cursor() as cursor:
            image_sql = f"""
                INSERT INTO {table_name} (vehicle_id, image_path, camera_id, capture_time) 
                VALUES (%s, %s, %s, %s)
            """
            image_values = []
            for item in image_data:
                image_path = item[0]
                vehicle_id = item[1]
                camera_id = item[2]
                capture_time = self._parse_capture_time(os.path.basename(image_path))
                image_values.append((vehicle_id, image_path, camera_id, capture_time))
            
            cursor.executemany(image_sql, image_values)
            print(f"成功插入 {len(image_values)} 条{table_name}记录")
        self.connection.commit()

    def _parse_capture_time(self, filename):
        """解析拍摄时间（Veri776无时间信息，返回None）"""
        return None

    def close(self):
        """关闭数据库连接"""
        if self.connection:
            self.connection.close()


def main():
    mysql_password = "root"
    dataset_path = str(Path(project_root).parent / "datas" / "VeRi-776")
    
    if not os.path.exists(dataset_path):
        print(f"错误：数据集路径 {dataset_path} 不存在")
        return
    
    db_manager = DatabaseManager(password=mysql_password)
    try:
        db_manager.import_veri776_dataset(dataset_path)
        print("Veri776数据集分表导入MySQL成功！")
    except Exception as e:
        print(f"导入失败: {e}")
        db_manager.connection.rollback()
    finally:
        db_manager.close()


if __name__ == "__main__":
    main()